from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options


def init_webdriver(
        driver_path: str = None,
        headless: bool = False,
        disable_gpu: bool = True,
        no_sandbox: bool = True,
        disable_dev_shm: bool = True,
) -> webdriver.Chrome:
    """
    初始化 Chrome WebDriver 的封装函数
    Args:
        driver_path (str): chromedriver 路径（默认 None，自动查找）
        headless (bool): 是否启用无头模式（默认 False）
        disable_gpu (bool): 是否禁用 GPU（默认 True）
        no_sandbox (bool): 是否禁用沙盒（默认 True）
        disable_dev_shm (bool): 是否禁用 /dev/shm（默认 True）
    Returns:
        webdriver.Chrome: 初始化后的 WebDriver 实例
    """
    # 配置 ChromeOptions
    chrome_options = Options()
    if headless:
        chrome_options.add_argument("--headless")  # 无头模式
    if disable_gpu:
        chrome_options.add_argument("--disable-gpu")  # 禁用 GPU 加速
    if no_sandbox:
        chrome_options.add_argument("--no-sandbox")  # 禁用沙盒模式
    if disable_dev_shm:
        chrome_options.add_argument("--disable-dev-shm-usage")  # 解决资源限制问题

    # 初始化 WebDriver
    if driver_path:
        service = Service(executable_path=driver_path)
    else:
        # 如果没有指定 driver_path，尝试自动查找（需安装 webdriver-manager）
        try:
            from webdriver_manager.chrome import ChromeDriverManager
            service = Service(ChromeDriverManager().install())
        except ImportError:
            raise ImportError("未安装 webdriver-manager，请手动指定 chromedriver 路径")

    driver = webdriver.Chrome(service=service, options=chrome_options)
    return driver
